from pysnmp.hlapi.asyncio import *
import asyncio

from src import get_raw_data,get_oid_data

raw_data = get_raw_data()
oid_data = get_oid_data()

policy = asyncio.WindowsSelectorEventLoopPolicy()
asyncio.set_event_loop_policy(policy)

async def _snmpwalk(ip,oid):
    community_data = CommunityData('public')
    transport_target = await UdpTransportTarget.create((ip, 161))  # SNMP Agent 地址和端口
    object_identity = ObjectIdentity(oid)

    result = {}

    async for error_indication, error_status, error_index, var_binds in walk_cmd(
        SnmpEngine(),
        community_data,
        transport_target,
        ContextData(),
        ObjectType(object_identity),
        lexicographicMode=False
    ):
        if error_indication:
            result = {"error": str(error_indication)+f"    oid:{oid},ip:{ip}"}
        elif error_status:
            result = {"error": f"{error_status.prettyPrint()} at {error_index}"+f"    oid:{oid},ip:{ip}"}
        else:
            for name, value in var_binds:
                result[str(name)] = str(value)
    # 等待所有数据收集完成后再返回
    return result

def snmpwalk(ip:str = None ,wr_type:str = None,mib_name:str = None):
    if ip == None or wr_type == None or mib_name == None: return {"error": "请输入正确的参数"}
    oid = raw_data[wr_type][mib_name]
    result = {}
    for key,value in asyncio.run(_snmpwalk(ip,oid)).items():
        if key == "error":
            result = {"error": value}
            print(result)
        else:
            result[oid_data[wr_type]["."+key]] = value
    return result

# 运行异步函数并等待返回结果
if __name__ == "__main__":
    result = snmpwalk("10.10.14.30","master","impCtrlMC0Ena")
    print(result)